import numpy as np
import qpsolvers as qp
import modern_robotics as mr

from ..robot import Robot


class RedundancyResolutionController:
    def __init__(self, robot: Robot):
        super().__init__()

        self._robot = robot
        self._Y = 0.01

    def ctrl(self, v_gripper_desired, v_base_desired):
        n = self._robot.dof

        t_err = np.sum(np.abs(v_gripper_desired[:3]))
        Y = 0.01

        Q = np.eye(n + 6)
        Q[:n, :n] *= Y
        Q[:2, :2] *= 1.0 / t_err
        Q[n:, n:] = (2.0 / t_err) * np.eye(6)

        v = v_gripper_desired

        q = self._robot.q
        Aeq = np.c_[self._robot.jacobe(q), np.eye(6)]
        beq = v.reshape((6,))

        Ain = np.zeros((n + 6, n + 6))
        bin = np.zeros(n + 6)

        ps = 0.1
        pi = 0.9

        Ain[:n, :n], bin[:n] = self._robot.joint_velocity_damper(q, ps, pi)
        c = np.concatenate(
            (np.zeros(2), -self._robot.jacobm(q, start=2).reshape((n - 2,)), np.zeros(6))
        )

        ke = 0.5
        bTe = self._robot.fkine(q, include_base=False)
        theta_e = np.arctan2(bTe[1, -1], bTe[0, -1])
        e = ke * theta_e

        lb = -np.r_[self._robot.qd_lim[: n], 10 * np.ones(6)]
        ub = np.r_[self._robot.qd_lim[: n], 10 * np.ones(6)]

        lb[0] = v_base_desired[2]
        lb[1] = v_base_desired[0]
        ub[0] = v_base_desired[2]
        ub[1] = v_base_desired[0]

        qd = qp.solve_qp(Q, c, Ain, bin, Aeq, beq, lb=lb, ub=ub, solver='cvxopt')
        qd = qd[: n]

        return qd
